// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use opentelemetry::trace::TracerProvider;
use opentelemetry_sdk::trace::SdkTracerProvider;
use tracing::Subscriber;
use tracing_subscriber::{EnvFilter, Layer, registry::LookupSpan};

/// The instrumentation scope name for the Google Cloud Rust client libraries.
/// This identifies spans generated by these libraries in the exported telemetry.
const INSTRUMENTATION_SCOPE_NAME: &str = "google-cloud-rust";

/// Creates a `tracing` layer that exports spans to the provided OpenTelemetry `SdkTracerProvider`.
///
/// The layer is filtered by the `RUST_LOG` environment variable (defaulting to "info").
///
/// # Examples
///
/// ```
/// use opentelemetry_sdk::trace::SdkTracerProvider;
/// use tracing_subscriber::{Registry, layer::SubscriberExt};
/// use integration_tests::observability::tracing::layer;
///
/// // Assume you have a configured SdkTracerProvider
/// let tracer_provider = SdkTracerProvider::builder().build();
///
/// let subscriber = Registry::default().with(layer(tracer_provider));
///
/// tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed");
/// ```
///
/// # Parameters
///
/// * `tracer_provider` - The configured OpenTelemetry `SdkTracerProvider`, which acts as the destination for spans.
///
/// # Type Parameters
///
/// * `S` - The type of the `tracing` Subscriber that this layer will be added to. For example,
///   `tracing_subscriber::Registry`.
pub fn layer<S>(tracer_provider: SdkTracerProvider) -> impl Layer<S>
where
    S: Subscriber + for<'span> LookupSpan<'span>,
{
    let tracer = tracer_provider.tracer(INSTRUMENTATION_SCOPE_NAME);
    let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));

    tracing_opentelemetry::layer()
        .with_tracer(tracer)
        .with_filter(env_filter)
}

#[cfg(test)]
mod tests {
    use super::*;
    use opentelemetry_sdk::error::OTelSdkError;
    use opentelemetry_sdk::trace::{SdkTracerProvider, SpanData, SpanExporter};
    use std::sync::{Arc, Mutex};
    use tracing_subscriber::{Registry, layer::SubscriberExt};

    #[derive(Debug, Clone)]
    struct InMemorySpanExporter {
        spans: Arc<Mutex<Vec<SpanData>>>,
    }

    impl InMemorySpanExporter {
        fn new() -> Self {
            Self {
                spans: Arc::new(Mutex::new(Vec::new())),
            }
        }
    }

    impl SpanExporter for InMemorySpanExporter {
        fn export(
            &self,
            batch: Vec<SpanData>,
        ) -> impl futures::Future<Output = Result<(), OTelSdkError>> + Send {
            let spans = self.spans.clone();
            async move {
                spans.lock().unwrap().extend(batch);
                Ok(())
            }
        }
    }

    #[tokio::test]
    async fn test_layer() {
        let exporter = InMemorySpanExporter::new();
        let spans = exporter.spans.clone();
        let provider = SdkTracerProvider::builder()
            .with_simple_exporter(exporter)
            .build();

        let layer = layer(provider.clone());
        let subscriber = Registry::default().with(layer);

        tracing::subscriber::with_default(subscriber, || {
            // Generate spans at different levels
            tracing::trace_span!("trace_span").in_scope(|| {});
            tracing::debug_span!("debug_span").in_scope(|| {});
            tracing::info_span!("info_span").in_scope(|| {});
            tracing::warn_span!("warn_span").in_scope(|| {});
            tracing::error_span!("error_span").in_scope(|| {});
        });

        // Ensure all spans are exported
        for _ in 0..10 {
            if let Err(e) = provider.force_flush() {
                eprintln!("Error flushing spans: {:?}", e);
                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
            } else {
                break;
            }
        }

        let exported_spans = spans.lock().unwrap();

        // Should only have info, warn, error (default filter is "info")
        let span_names: Vec<&str> = exported_spans.iter().map(|s| s.name.as_ref()).collect();
        assert!(
            span_names.contains(&"info_span"),
            "Expected 'info_span' to be present in: {:?}",
            span_names
        );
        assert!(
            span_names.contains(&"warn_span"),
            "Expected 'warn_span' to be present in: {:?}",
            span_names
        );
        assert!(
            span_names.contains(&"error_span"),
            "Expected 'error_span' to be present in: {:?}",
            span_names
        );
        assert!(
            !span_names.contains(&"trace_span"),
            "Did not expect 'trace_span' to be present in: {:?}",
            span_names
        );
        assert!(
            !span_names.contains(&"debug_span"),
            "Did not expect 'debug_span' to be present in: {:?}",
            span_names
        );
    }
}
